#define STB_IMAGE_IMPLEMENTATION

#include "LumineTexture.hpp"
#include "Core/Context/LumineContext.hpp"
#include "Core/Buffer/LumineBuffer.hpp"
#include "Core/Buffer/ImageBuffer.hpp"
#include "Core/Buffer/StagingBuffer.hpp"
#include "Utils/LumineUtils.hpp"

void Lumine::LumineTexture::createImage(uint32_t width, uint32_t height, vk::Format format, vk::ImageTiling tiling, vk::ImageUsageFlags usage, vk::MemoryPropertyFlags properties, vk::raii::Image &image, vk::raii::DeviceMemory &imageMemory)
{
    vk::ImageCreateInfo imageInfo;
    imageInfo.setImageType(vk::ImageType::e2D)
        .setFormat(format)
        .setExtent({width, height, 1})
        .setMipLevels(1)
        .setArrayLayers(1)
        .setSamples(vk::SampleCountFlagBits::e1)
        .setTiling(tiling)
        .setUsage(usage)
        .setSharingMode(vk::SharingMode::eExclusive);
    image = vk::raii::Image(LumineContext::Get().Device(), imageInfo);

    vk::MemoryRequirements memReq = image.getMemoryRequirements();
    vk::MemoryAllocateInfo allocInfo;
    allocInfo.setAllocationSize(memReq.size)
        .setMemoryTypeIndex(LumineBuffer::findMemoryType(memReq.memoryTypeBits, properties));

    imageMemory = vk::raii::DeviceMemory(LumineContext::Get().Device(), allocInfo);
    image.bindMemory(imageMemory, 0);
}

stbi_uc *Lumine::LumineTexture::loadPixels(const std::string &path, int &w, int &h, vk::DeviceSize size)
{
    int channels;
    // 强制加载为 4 个通道 (RGBA)
    stbi_uc *pixels = stbi_load(path.c_str(), &w, &h, &channels, STBI_rgb_alpha);
    size = w * h * 4;
    if (!pixels)
    {
        throw std::runtime_error("Failed to load texture image!");
    }
    return pixels;
}

void Lumine::LumineTexture::transitionImageLayout(const vk::raii::Image &image,vk::ImageAspectFlags aspect, vk::ImageLayout oldLayout, vk::ImageLayout newLayout)
{
    // 获取上下文中的队列族索引
    auto &context = LumineContext::Get();
    uint32_t transferIndex = context.TransferIndex();
    uint32_t graphicIndex = context.GraphicIndex();

    Utils::execSingleTimeCmds([&](const vk::raii::CommandBuffer &cmdBuffer)
                              {
        vk::ImageMemoryBarrier barrier{};
        barrier.setOldLayout(oldLayout)
               .setNewLayout(newLayout)
               .setImage(*image)
               .setSubresourceRange({aspect, 0, 1, 0, 1}); 

        vk::PipelineStageFlags sourceStage;
        vk::PipelineStageFlags destinationStage;

        if (oldLayout == vk::ImageLayout::eUndefined && newLayout == vk::ImageLayout::eTransferDstOptimal)
        {
            // TRANSITION 1: UNDEFINED -> TRANSFER_DST (准备写入)
            barrier.setSrcAccessMask(vk::AccessFlagBits::eNone);
            barrier.setDstAccessMask(vk::AccessFlagBits::eTransferWrite);
            sourceStage = vk::PipelineStageFlagBits::eTopOfPipe;
            destinationStage = vk::PipelineStageFlagBits::eTransfer;
            
            // 初始转换不需要所有权转移
            barrier.setSrcQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED);
            barrier.setDstQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED);
        }
        else if (oldLayout == vk::ImageLayout::eTransferDstOptimal && newLayout == vk::ImageLayout::eShaderReadOnlyOptimal)
        {
            // TRANSITION 2: TRANSFER_DST -> SHADER_READ_ONLY (准备读取)
            barrier.setSrcAccessMask(vk::AccessFlagBits::eTransferWrite);
            barrier.setDstAccessMask(vk::AccessFlagBits::eShaderRead);
            sourceStage = vk::PipelineStageFlagBits::eTransfer;
            destinationStage = vk::PipelineStageFlagBits::eFragmentShader;
        } 
        else if(oldLayout == vk::ImageLayout::eUndefined && newLayout == vk::ImageLayout::eDepthStencilAttachmentOptimal){
            barrier.srcAccessMask = vk::AccessFlagBits::eNone;
            barrier.dstAccessMask = vk::AccessFlagBits::eDepthStencilAttachmentRead |vk::AccessFlagBits::eDepthStencilAttachmentWrite;
            sourceStage = vk::PipelineStageFlagBits::eTopOfPipe;
            destinationStage = vk::PipelineStageFlagBits::eEarlyFragmentTests;
        }
        cmdBuffer.pipelineBarrier( sourceStage, destinationStage, {}, {}, nullptr, barrier ); });
}

void Lumine::LumineTexture::copyBufferToImage(const vk::raii::Buffer &buffer, vk::raii::Image &image, uint32_t width, uint32_t height)
{
    Utils::execSingleTimeCmds([&](const vk::raii::CommandBuffer &cmdBuffer)
                              {
    vk::BufferImageCopy region;
    vk::ImageSubresourceLayers subresource = {vk::ImageAspectFlagBits::eColor, 0, 0, 1};
    region.setBufferOffset(0)
        .setBufferRowLength(0)
        .setBufferImageHeight(0)
        .setImageSubresource(subresource)
        .setImageOffset({0,0,0})
        .setImageExtent({width,height,1}); 
    cmdBuffer.copyBufferToImage(buffer,image,vk::ImageLayout::eTransferDstOptimal,{region}); });
}


